from PyQt6.QtWidgets import QMainWindow, QPushButton, QMenu, QApplication, QMessageBox
from PyQt6.QtCore import Qt, QPoint
from main_window import AIMainWindow
from styles import FLOAT_BUTTON_STYLE

class FloatingAIAssistant(QMainWindow):
    def __init__(self):
        super().__init__()
        self.drag_start_global_pos = None
        self.ai_main_window = None
        self.version = "1.3.0"  # 更新版本号
        self.last_model = None  # 保存上次使用的模型
        self.last_conversation_id = None  # 保存上次使用的对话ID
        self.initUI()
        
    def initUI(self):
        self.setWindowTitle("AI 助手")
        # 设置窗口属性
        self.setWindowFlags(
            Qt.WindowType.FramelessWindowHint | 
            Qt.WindowType.WindowStaysOnTopHint |
            Qt.WindowType.Tool
        )
        self.setAttribute(Qt.WidgetAttribute.WA_TranslucentBackground)
        
        # 创建按钮
        self.toggle_button = QPushButton("AI", self)
        self.toggle_button.setFixedSize(50, 50)
        self.toggle_button.setStyleSheet(FLOAT_BUTTON_STYLE)
        
        # 设置窗口大小和初始位置
        self.setFixedSize(50, 50)
        screen = QApplication.primaryScreen().geometry()
        self.move(screen.width() - 100, screen.height() // 2)
        
        # 绑定鼠标事件
        self.toggle_button.mousePressEvent = self.button_mouse_press
        self.toggle_button.mouseMoveEvent = self.button_mouse_move
        self.toggle_button.mouseReleaseEvent = self.button_mouse_release
        
        # 创建右键菜单
        self.toggle_button.setContextMenuPolicy(Qt.ContextMenuPolicy.CustomContextMenu)
        self.toggle_button.customContextMenuRequested.connect(self.show_context_menu)
        
        # 更新悬浮按钮的提示文本
        self.setToolTip(f"AI助手 v{self.version}\n点击展开/收起主界面\n按住拖动调整位置")
        
        self.show()
        
    def button_mouse_press(self, event):
        if event.button() == Qt.MouseButton.LeftButton:
            self.drag_start_global_pos = event.globalPosition().toPoint()
    
    def button_mouse_move(self, event):
        if not (event.buttons() & Qt.MouseButton.LeftButton) or not self.drag_start_global_pos:
            return
            
        current_global_pos = event.globalPosition().toPoint()
        delta = current_global_pos - self.drag_start_global_pos
        new_pos = self.pos() + delta
        
        screen = QApplication.primaryScreen().geometry()
        new_pos.setX(max(0, min(new_pos.x(), screen.width() - self.width())))
        new_pos.setY(max(0, min(new_pos.y(), screen.height() - self.height())))
        
        self.move(new_pos)
        self.drag_start_global_pos = current_global_pos
    
    def button_mouse_release(self, event):
        if event.button() == Qt.MouseButton.LeftButton:
            self.toggle_ai_main_window()
            self.drag_start_global_pos = None

    def toggle_ai_main_window(self):
        if self.ai_main_window is None:
            self.ai_main_window = AIMainWindow(self)
            
            # 恢复上次的模型和对话状态
            if self.last_model:
                # 设置模型
                index = self.ai_main_window.model_selector.findText(self.last_model)
                if index >= 0:
                    self.ai_main_window.model_selector.setCurrentIndex(index)
                    
                # 设置对话
                if self.last_conversation_id is not None:
                    self.ai_main_window.conversation_manager.set_current_conversation(self.last_model, self.last_conversation_id)
                    self.ai_main_window.update_conversation_selector(self.last_model, self.last_conversation_id)
            
            # 获取屏幕大小
            screen = QApplication.primaryScreen().geometry()
            
            # 获取按钮的全局位置
            button_global_pos = self.mapToGlobal(QPoint(0, 0))
            
            # 计算主界面位置
            main_window_x = button_global_pos.x()
            main_window_y = button_global_pos.y() + self.height()
            
            # 调整位置确保窗口在屏幕内
            if main_window_x + self.ai_main_window.width() > screen.width():
                main_window_x = button_global_pos.x() + self.width() - self.ai_main_window.width() - 25
            
            if main_window_y + self.ai_main_window.height() > screen.height():
                main_window_y = button_global_pos.y() - self.ai_main_window.height()
            
            if main_window_y < 0:
                main_window_y = button_global_pos.y() + self.height()
                if main_window_y + self.ai_main_window.height() > screen.height():
                    main_window_y = 10
            
            main_window_x = max(0, main_window_x)
            
            self.ai_main_window.move(main_window_x, main_window_y)
            self.hide()
            self.ai_main_window.show()
        else:
            self.ai_main_window.close()
            self.ai_main_window = None
            self.show()

    def show_context_menu(self, position):
        context_menu = QMenu(self)
        
        # 只保留"退出"选项
        exit_action = context_menu.addAction("退出")
        exit_action.triggered.connect(QApplication.quit)
        
        context_menu.exec(self.toggle_button.mapToGlobal(position))